#!/bin/bash -e

wait_for_containerd_ready()
{
  until ctr --connect-timeout 1s version&>/dev/null
  do
      sleep 1
  done
}

preload_rke2_images()
{
    readonly RKE2_IMAGES_DIR="/var/lib/rancher/rke2/agent/images"
    readonly TMP_IMAGES_DIR="/var/lib/rancher/tmp/images"
    readonly IMAGES_LISTS_DIR="/tmp/images-lists"
    mkdir -p $RKE2_IMAGES_DIR
    mkdir -p $TMP_IMAGES_DIR
    mkdir -p $IMAGES_LISTS_DIR

    # If the installation mode is ISO, use rsync instead of bind-mount to mitigate the potential installation failure due to slow media
    # Otherwise (PXE), use bind-mount directly to speed up the process
    echo "Copying RKE2 images to the target location..."
    cp /root/images/harvester/images/rke2-images.*.tar.zst $RKE2_IMAGES_DIR
    echo "Copying remaining images temporary location..."
    #cp -r ${ISOMNT}/bundle/harvester/images/ $TARGET/$TMP_IMAGES_DIR
    cp -r /root/images/harvester/images/* $TMP_IMAGES_DIR
    #cp -r ${ISOMNT}/bundle/harvester/images-lists/ $TARGET/$IMAGES_LISTS_DIR
    cp -r /root/images/harvester/images-lists/* $IMAGES_LISTS_DIR

    echo "Loading images. This may take a few minutes..."
    install_mode="$INSTALL_MODE"
    tmp_images_dir="$TMP_IMAGES_DIR"
    images_lists_dir="$IMAGES_LISTS_DIR"
    set -e

    mkdir -p /tmp/local
    inst_tmp=$(mktemp -d -p /tmp/local)
    trap "rm -rf $inst_tmp" exit

    # extract RKE2 tarball from image
    image_list=$(ls /var/lib/rancher/agent/images/rancherd-bootstrap-images-*.txt | head -n 1)
    if [ -z "$image_list" ]; then
      echo "[ERROR] Fail to get rancherd bootstrap images list."
      exit 1
    fi
    rke2_image=$(grep 'docker.io/rancher/system-agent-installer-rke2:' $image_list)
    wharfie --images-dir /var/lib/rancher/agent/images/ $rke2_image $inst_tmp

    # extract RKE2 binary
    rke2_tmp="$inst_tmp/rke2"
    if [ ${HOSTTYPE} = "aarch64" ]
    then
      mkdir -p $rke2_tmp && tar xf $inst_tmp/rke2.linux-arm64.tar.gz -C $rke2_tmp
    else
      mkdir -p $rke2_tmp && tar xf $inst_tmp/rke2.linux-amd64.tar.gz -C $rke2_tmp
    fi

    $rke2_tmp/bin/rke2 server &> /rke2.log &

    export PATH=/var/lib/rancher/rke2/bin:$PATH
    export CONTAINERD_ADDRESS=/run/k3s/containerd/containerd.sock

    wait_for_containerd_ready

    echo "Stop RKE2 and remove temporary RKE2 files..."
    pkill rke2
    rm -f /rke2.log
    rm -f /etc/rancher/rke2/rke2.yaml
    rm -rf /var/lib/rancher/rke2/server
    rm -rf /var/lib/rancher/rke2/agent/pod-manifests/*

    echo "Start containerd..."
    containerd -c /var/lib/rancher/rke2/agent/etc/containerd/config.toml -a $CONTAINERD_ADDRESS --state /run/k3s/containerd --root /var/lib/rancher/rke2/agent/containerd &> /containerd.log &

    wait_for_containerd_ready

    # load images
    for i in $tmp_images_dir/*.tar.zst; do
        echo "Load images from $i"
        zstd -d $i -o /usr/local/images.tar
        ctr -n k8s.io images import --no-unpack /usr/local/images.tar
        rm /usr/local/images.tar
    done

    # make sure all preloading images are ready
    for i in $images_lists_dir/*.txt; do
      stdbuf -oL ./ctr-check-images.sh $i
    done

    # tearing down containerd/RKE2
    echo "Stop containerd..."
    pkill containerd
    rm -f /containerd.log
    rm -rf ${RKE2_IMAGES_DIR}/*
    rm -rf ${TMP_IMAGES_DIR}
    rm -rf ${IMAGES_LISTS_DIR}/*
}

preload_rancherd_images()
{
    mkdir -p /var/lib/rancher/agent/images
    cp /root/images/rancherd/images/* /var/lib/rancher/agent/images
}


do_preload()
{
    # Preload all images now
    preload_rancherd_images
    preload_rke2_images

}

do_install_wharfie()
{

    # Check if wharfie exists
    if ! command -v wharfie &> /dev/null
    then
        echo "wharfie not found, starting download..."
        
        # Set architecture variable
        arch=""
        case "$(uname -m)" in
            arm64|aarch64)
                arch="arm64"
                ;;
            x86_64|amd64)
                arch="amd64"
                ;;
            *)
                echo "Unsupported architecture: $(uname -m)"
                exit 1
                ;;
        esac
        
        # Download wharfie
        curl -L -o wharfie https://ghproxy.cn/https://github.com/rancher/wharfie/releases/download/v0.6.8/wharfie-${arch}
        
        # Add execute permission
        chmod +x wharfie
        
        # Move to /usr/local/bin
        sudo mv wharfie /usr/local/bin/wharfie
        
        echo "wharfie has been successfully installed."
    else
        echo "wharfie already exists."
    fi

}


do_install_wharfie
# Preload images
do_preload


